package com.kssoft.ebd.spider.http;

import com.fasterxml.jackson.databind.ObjectMapper;
import org.apache.http.NameValuePair;
import org.apache.http.client.entity.UrlEncodedFormEntity;
import org.apache.http.client.methods.HttpPost;
import org.apache.http.entity.ContentType;
import org.apache.http.entity.StringEntity;
import org.apache.http.entity.mime.MultipartEntityBuilder;
import org.apache.http.entity.mime.content.FileBody;
import org.apache.http.entity.mime.content.StringBody;
import org.apache.http.message.BasicNameValuePair;
import us.codecraft.webmagic.Request;
import us.codecraft.webmagic.Site;
import us.codecraft.webmagic.downloader.HttpClientRequestContext;
import us.codecraft.webmagic.downloader.HttpUriRequestConverter;
import us.codecraft.webmagic.proxy.Proxy;

import java.io.File;
import java.nio.charset.StandardCharsets;
import java.util.*;
import java.util.stream.Collectors;

public class CustomHttpUriRequestConverter extends HttpUriRequestConverter {

    private static final ObjectMapper objectMapper = new ObjectMapper();

    @Override
    public HttpClientRequestContext convert(Request request, Site site, Proxy proxy) {
        String method = getStringExtra(request, "method", "GET").toUpperCase();
        if (!"POST".equals(method)) {
            return super.convert(request, site, proxy);
        }

        String url = request.getUrl();
        Map<String, Object> params = getMapExtra(request, "params");

        HttpPost post = new HttpPost(url);
        String contentType = getStringExtra(request, "contentType", "form").toLowerCase();

        switch (contentType) {
            case "json":
                Object jsonBody = request.getExtra("jsonBody");
                String jsonStr;
                try {
                    if (jsonBody instanceof String) {
                        jsonStr = (String) jsonBody;
                    } else {
                        jsonStr = objectMapper.writeValueAsString(jsonBody);
                    }
                } catch (Exception e) {
                    throw new RuntimeException("Failed to serialize JSON", e);
                }
                post.setEntity(new StringEntity(jsonStr, ContentType.APPLICATION_JSON));
                break;

            case "multipart":
                Map<String, Object> multipart = getMapExtra(request, "multipart");
                if (multipart != null) {
                    MultipartEntityBuilder builder = MultipartEntityBuilder.create();
                    for (Map.Entry<String, Object> entry : multipart.entrySet()) {
                        String key = entry.getKey();
                        Object val = entry.getValue();
                        if (val instanceof File) {
                            builder.addPart(key, new FileBody((File) val));
                        } else {
                            builder.addPart(key, new StringBody(String.valueOf(val), ContentType.TEXT_PLAIN));
                        }
                    }
                    post.setEntity(builder.build());
                }
                break;

            case "form":
            default:
                if (params != null && !params.isEmpty()) {
                    List<NameValuePair> pairs = params.entrySet().stream()
                            .map(e -> new BasicNameValuePair(e.getKey(), e.getValue().toString()))
                            .collect(Collectors.toList());
                    post.setEntity(new UrlEncodedFormEntity(pairs, StandardCharsets.UTF_8));
                }
                break;
        }
        HttpClientRequestContext convert = super.convert(request, site, proxy);
        convert.setHttpUriRequest(post);
        return convert;
    }

    private String getStringExtra(Request request, String key, String defaultValue) {
        Object value = request.getExtra(key);
        return (value instanceof String) ? (String) value : defaultValue;
    }

    @SuppressWarnings("unchecked")
    private Map<String, Object> getMapExtra(Request request, String key) {
        Object value = request.getExtra(key);
        return (value instanceof Map) ? (Map<String, Object>) value : null;
    }
}
